package com.mini.grpc.common;

import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Callable;
import java.util.concurrent.CancellationException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Context提供了一种在API边界之间传递截止时间、取消信号和其他请求范围值的方法。
 */
public class Context {
    
    /**
     * 根上下文，所有其他上下文都是从它派生的
     */
    public static final Context ROOT = new Context();
    
    private static final ThreadLocal<Context> CURRENT_CONTEXT = new ThreadLocal<Context>() {
        @Override
        protected Context initialValue() {
            return ROOT;
        }
    };
    
    private final Context parent;
    private final Map<Key<?>, Object> keyValueEntries;
    private final CancellationListener cancellationListener;
    private final Deadline deadline;
    private final AtomicReference<CancellationCause> cancellationCause;
    
    private Context() {
        this.parent = null;
        this.keyValueEntries = new HashMap<>();
        this.cancellationListener = null;
        this.deadline = null;
        this.cancellationCause = new AtomicReference<>();
    }
    
    private Context(Context parent, Map<Key<?>, Object> keyValueEntries, CancellationListener cancellationListener,
                   Deadline deadline) {
        this.parent = parent;
        this.keyValueEntries = keyValueEntries;
        this.cancellationListener = cancellationListener;
        this.deadline = deadline;
        this.cancellationCause = new AtomicReference<>();
    }
    
    /**
     * 获取当前线程的上下文
     *
     * @return 当前上下文
     */
    public static Context current() {
        return CURRENT_CONTEXT.get();
    }
    
    /**
     * 在指定的上下文中执行任务
     *
     * @param context 要使用的上下文
     * @param runnable 要执行的任务
     */
    public static void run(Context context, Runnable runnable) {
        Context previous = CURRENT_CONTEXT.get();
        CURRENT_CONTEXT.set(context);
        try {
            runnable.run();
        } finally {
            CURRENT_CONTEXT.set(previous);
        }
    }
    
    /**
     * 在指定的上下文中执行可调用任务
     *
     * @param context 要使用的上下文
     * @param callable 要执行的可调用任务
     * @param <V> 返回值类型
     * @return 任务的返回值
     * @throws Exception 如果任务抛出异常
     */
    public static <V> V call(Context context, Callable<V> callable) throws Exception {
        Context previous = CURRENT_CONTEXT.get();
        CURRENT_CONTEXT.set(context);
        try {
            return callable.call();
        } finally {
            CURRENT_CONTEXT.set(previous);
        }
    }
    
    /**
     * 获取上下文中的值
     *
     * @param key 键
     * @param <V> 值类型
     * @return 值，如果不存在则返回null
     */
    @SuppressWarnings("unchecked")
    public <V> V get(Key<V> key) {
        Object value = keyValueEntries.get(key);
        if (value != null) {
            return (V) value;
        }
        if (parent != null) {
            return parent.get(key);
        }
        return null;
    }
    
    /**
     * 创建一个带有键值对的新上下文
     *
     * @param key 键
     * @param value 值
     * @param <V> 值类型
     * @return 新的上下文
     */
    public <V> Context withValue(Key<V> key, V value) {
        Map<Key<?>, Object> newKeyValueEntries = new HashMap<>(keyValueEntries);
        newKeyValueEntries.put(key, value);
        return new Context(this, newKeyValueEntries, cancellationListener, deadline);
    }
    
    /**
     * 创建一个带有截止时间的新上下文
     *
     * @param deadline 截止时间
     * @return 新的上下文
     */
    public Context withDeadline(Deadline deadline) {
        return new Context(this, keyValueEntries, cancellationListener, deadline);
    }
    
    /**
     * 创建一个带有超时的新上下文
     *
     * @param duration 超时时间
     * @param unit 时间单位
     * @return 新的上下文
     */
    public Context withTimeout(long duration, TimeUnit unit) {
        return withDeadline(Deadline.after(duration, unit));
    }
    
    /**
     * 创建一个可取消的新上下文
     *
     * @param cancellationListener 取消监听器
     * @return 新的上下文
     */
    public Context withCancellation(CancellationListener cancellationListener) {
        return new Context(this, keyValueEntries, cancellationListener, deadline);
    }
    
    /**
     * 检查上下文是否已取消
     *
     * @return 如果上下文已取消则返回true，否则返回false
     */
    public boolean isCancelled() {
        return cancellationCause.get() != null || (parent != null && parent.isCancelled());
    }
    
    /**
     * 取消上下文
     *
     * @param cause 取消原因
     * @return 如果成功取消则返回true，否则返回false
     */
    public boolean cancel(CancellationCause cause) {
        if (cancellationCause.compareAndSet(null, cause)) {
            if (cancellationListener != null) {
                cancellationListener.cancelled(this);
            }
            return true;
        }
        return false;
    }
    
    /**
     * 获取取消原因
     *
     * @return 取消原因，如果未取消则返回null
     */
    public CancellationCause cancellationCause() {
        CancellationCause cause = cancellationCause.get();
        if (cause != null) {
            return cause;
        }
        if (parent != null) {
            return parent.cancellationCause();
        }
        return null;
    }
    
    /**
     * 获取截止时间
     *
     * @return 截止时间，如果没有则返回null
     */
    public Deadline getDeadline() {
        if (deadline != null) {
            return deadline;
        }
        if (parent != null) {
            return parent.getDeadline();
        }
        return null;
    }
    
    /**
     * 检查是否已超过截止时间
     *
     * @return 如果已超过截止时间则返回true，否则返回false
     */
    public boolean isDeadlineExceeded() {
        Deadline deadline = getDeadline();
        return deadline != null && deadline.isExpired();
    }
    
    /**
     * 上下文键，用于存储和检索值
     *
     * @param <V> 值类型
     */
    public static final class Key<V> {
        private final String name;
        
        private Key(String name) {
            this.name = name;
        }
        
        /**
         * 创建一个新的键
         *
         * @param name 键名
         * @param <V> 值类型
         * @return 新的键
         */
        public static <V> Key<V> of(String name) {
            return new Key<>(name);
        }
        
        @Override
        public boolean equals(Object obj) {
            if (this == obj) {
                return true;
            }
            if (obj == null || getClass() != obj.getClass()) {
                return false;
            }
            Key<?> other = (Key<?>) obj;
            return name.equals(other.name);
        }
        
        @Override
        public int hashCode() {
            return name.hashCode();
        }
        
        @Override
        public String toString() {
            return "Key{name=" + name + "}";
        }
    }
    
    /**
     * 截止时间，表示操作必须在指定时间之前完成
     */
    public static final class Deadline {
        private final long deadlineNanos;
        
        private Deadline(long deadlineNanos) {
            this.deadlineNanos = deadlineNanos;
        }
        
        /**
         * 创建一个相对于当前时间的截止时间
         *
         * @param duration 时间长度
         * @param unit 时间单位
         * @return 新的截止时间
         */
        public static Deadline after(long duration, TimeUnit unit) {
            return new Deadline(System.nanoTime() + unit.toNanos(duration));
        }
        
        /**
         * 检查截止时间是否已过期
         *
         * @return 如果已过期则返回true，否则返回false
         */
        public boolean isExpired() {
            return deadlineNanos <= System.nanoTime();
        }
        
        /**
         * 获取剩余时间
         *
         * @param unit 时间单位
         * @return 剩余时间
         * @throws TimeoutException 如果截止时间已过期
         */
        public long timeRemaining(TimeUnit unit) throws TimeoutException {
            long remainingNanos = deadlineNanos - System.nanoTime();
            if (remainingNanos <= 0) {
                throw new TimeoutException("Deadline exceeded");
            }
            return unit.convert(remainingNanos, TimeUnit.NANOSECONDS);
        }
    }
    
    /**
     * 取消原因，表示上下文被取消的原因
     */
    public static final class CancellationCause {
        private final String reason;
        private final Throwable cause;
        
        /**
         * 创建一个取消原因
         *
         * @param reason 取消原因描述
         */
        public CancellationCause(String reason) {
            this(reason, null);
        }
        
        /**
         * 创建一个带有异常的取消原因
         *
         * @param reason 取消原因描述
         * @param cause 导致取消的异常
         */
        public CancellationCause(String reason, Throwable cause) {
            this.reason = reason;
            this.cause = cause;
        }
        
        /**
         * 获取取消原因描述
         *
         * @return 取消原因描述
         */
        public String getReason() {
            return reason;
        }
        
        /**
         * 获取导致取消的异常
         *
         * @return 导致取消的异常，可能为null
         */
        public Throwable getCause() {
            return cause;
        }
        
        /**
         * 创建一个CancellationException
         *
         * @return 新的CancellationException
         */
        public CancellationException asException() {
            CancellationException exception = new CancellationException(reason);
            if (cause != null) {
                exception.initCause(cause);
            }
            return exception;
        }
    }
    
    /**
     * 取消监听器，用于在上下文被取消时接收通知
     */
    public interface CancellationListener {
        /**
         * 当上下文被取消时调用
         *
         * @param context 被取消的上下文
         */
        void cancelled(Context context);
    }
    
    /**
     * 创建一个可以在指定执行器中执行任务的上下文包装器
     *
     * @param executor 执行器
     * @return 上下文包装器
     */
    public Context.ContextExecutor withExecutor(Executor executor) {
        return new ContextExecutor(this, executor);
    }
    
    /**
     * 上下文执行器，用于在指定的上下文中执行任务
     */
    public static final class ContextExecutor implements Executor {
        private final Context context;
        private final Executor executor;
        
        private ContextExecutor(Context context, Executor executor) {
            this.context = context;
            this.executor = executor;
        }
        
        @Override
        public void execute(Runnable command) {
            executor.execute(() -> Context.run(context, command));
        }
    }
} 